#!/bin/bash

# Copyright 2014-2022 Security Onion Solutions, LLC

#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

cd "$(dirname "$0")" || exit 255

source ../salt/common/tools/sbin/so-common
source ./so-functions

script_run="$1"

retry_count=10
retry_sleep=5
warning_prefix="[WARNING]"
info_prefix="[INFO   ]"
error_prefix="[ERROR  ]"

if [[ $script_run == true ]]; then
	preflight_log="${2:-'/root/preflight.log'}"
else
	preflight_log='/root/preflight.log'
fi

check_default_repos() {
	local ret_code=0
	local repo_str='  Checking OS default repos with '
	if [[ $script_run == true ]]; then
		printf '%s' "$repo_str"
	else
		printf '%s' "$repo_str" | tee -a "$preflight_log"
	fi

	if [[ $OS == 'centos' ]]; then
		if [[ $script_run == true ]]; then
			printf '%s' 'yum update.'
		else
			printf '%s' 'yum update.' | tee -a "$preflight_log"
		fi
		echo "" >> "$preflight_log"
		yum -y check-update >> $preflight_log 2>&1
		ret_code=$?
		if [[ $ret_code == 0 || $ret_code == 100 ]]; then
			printf '%s\n' '  SUCCESS'
			ret_code=0
		else
			printf '%s\n' '  FAILURE'
		fi
	else
		if [[ $script_run == true ]]; then
			printf '%s' 'apt update.'
		else
			printf '%s' 'apt update.' | tee -a "$preflight_log"
		fi
		echo "" >> "$preflight_log"
		retry 50 10 "apt-get -y update" >> $preflight_log 2>&1
		ret_code=$?
		[[ $ret_code == 0 ]] && printf '%s\n' '  SUCCESS' || printf '%s\n' '  FAILURE'

	fi

	return $ret_code
}

check_new_repos() {
	local repo_url_str='  Checking repo URLs added by setup.'
	if [[ $script_run == true ]]; then
		printf '%s' "$repo_url_str"
	else
		printf '%s' "$repo_url_str" | tee -a "$preflight_log"
	fi

	if [[ $OS == 'centos' ]]; then
		local repo_arr=( 
			"https://download.docker.com/linux/centos/docker-ce.repo"
			"https://repo.securityonion.net/file/securityonion-repo/keys/SALTSTACK-GPG-KEY.pub"
			"https://download.docker.com/linux/ubuntu/gpg"
			"https://packages.wazuh.com/key/GPG-KEY-WAZUH"
			"https://packages.wazuh.com/3.x/yum/"
			)
	else
		local ubuntu_version
		ubuntu_version=$(grep VERSION_ID /etc/os-release 2> /dev/null | awk -F '[ "]' '{print $2}')
		local repo_arr=(
			"https://download.docker.com/linux/ubuntu/gpg"
			"https://download.docker.com/linux/ubuntu"
			"https://repo.securityonion.net/file/securityonion-repo/ubuntu/$ubuntu_version/amd64/salt/SALTSTACK-GPG-KEY.pub"
			"https://packages.wazuh.com/key/GPG-KEY-WAZUH"
			"https://packages.wazuh.com"
		)
	fi

	__check_url_arr "${repo_arr[@]}"
	local ret_code=$?
	[[ $ret_code == 0 ]] && printf '%s\n' '  SUCCESS' || printf '%s\n' '  FAILURE'
	return $ret_code
}

check_misc_urls() {
	local misc_url_str='  Checking various other URLs used by setup.'
	if [[ $script_run == true ]]; then
		printf '%s' "$misc_url_str"
	else
		printf '%s' "$misc_url_str" | tee -a "$preflight_log"
	fi

	local so_version
	so_version=$(cat ../VERSION)
	local url_arr=(
		"https://raw.githubusercontent.com/Security-Onion-Solutions/securityonion/master/KEYS"
		"https://github.com/Neo23x0/signature-base"
		"https://sigs.securityonion.net/$so_version/securityonion-$so_version.iso.sig"
		"https://ghcr.io/"
		"https://rules.emergingthreats.net/open/"
		"https://rules.emergingthreatspro.com/"
	)

	__check_url_arr "${url_arr[@]}"
	local ret_code=$?
	[[ $ret_code == 0 ]] && printf '%s\n' '  SUCCESS' || printf '%s\n' '  FAILURE'
	return $ret_code
}

__check_url_arr() {
	local ret_code=0
	echo "" >> "$preflight_log"
	for url in "$@"; do
		# Reset vars
		local status=999 # Set status to something outside the range of normal HTTP codes but above the 200 range
		local ret=1
		local count=0

		while [[ $ret != 0 && $count -lt $retry_count ]]; do
			((count++))
			[[ $count != 1 ]] && sleep $retry_sleep
			status=$(curl -s -o /dev/null -w "%{http_code}" -L "$url" 2> /dev/null)
			ret=$?
			local count_str
			printf -v count_str '%02d' "$count"
			[[ $ret != 0 ]] && echo "$warning_prefix ($count_str/$retry_count) Could not reach $url, curl error code: $ret" >> "$preflight_log"
		done

		if [[ $ret == 0 ]]; then
			url_success_str="Successfully reached $url"
			if [[ $status -ge 400 ]]; then
				echo "$warning_prefix $url_success_str but server responded with HTTP code $status." >> "$preflight_log"
			else
				printf '%s\n' "$info_prefix $url_success_str" >> "$preflight_log"
			fi
		else
			ret_code=1
			echo "$error_prefix Could not reach $url after $retry_count attempts." >> "$preflight_log"
		fi
	done
	echo "" >> "$preflight_log"
	return $ret_code
}

preflight_prereqs() {
	local ret_code=0
	
	if [[ $OS == 'centos' ]]; then
		: # no-op to match structure of other checks for $OS var
	else
		retry 50 10 "apt-get -y install curl" >> "$preflight_log" 2>&1 || ret_code=1
	fi

	return $ret_code
}

main() {
	local intro_str="Beginning pre-flight checks."
	local success_str="Pre-flight checks completed successfully!"
	local fail_str="Pre-flight checks could not complete."

	[[ -f $preflight_log ]] || touch "$preflight_log"

	detect_os "$preflight_log"

	if [[ $script_run == true ]]; then
		echo "$intro_str"
	else
		echo "$intro_str" | tee "$preflight_log"
	fi

	check_default_repos &&\
	preflight_prereqs &&\
	check_new_repos &&\
	check_misc_urls

	local success=$?
	
	echo ""
	if [[ $success == 0 ]]; then
		if [[ $script_run == true ]]; then
			echo "$success_str"
		else
			echo "$success_str" | tee -a "$preflight_log"
			echo ""
		fi
	else
		if [[ $script_run == true ]]; then
			echo "$fail_str"
		else
			echo "$fail_str" | tee -a "$preflight_log"
			echo "Check $preflight_log for details."
			echo ""
		fi
	fi

	exit $success
}

main
